from typing import Optional, Callable

import os
import torch
import torch.utils.data as Data
from PIL import Image


class imageDataset(Data.Dataset):
    '''
    在训练中，2D训练数据加载类
    '''

    def __init__(self, path, batch_size=1, step_size=1, transform: Optional[Callable] = None):
        '''
        初始化过程
        :param path:图像文件的路径
        :param batch_size: 批次
        :param step_size: 步长
        :param transform: 变换
        '''
        self.path = path
        self.batch_size = batch_size
        self.step_size = step_size
        self.transform = transform

        length = 0
        first_img = True
        for img_f in os.listdir(self.path):
            if os.path.splitext(img_f)[-1] == ".jpg":
                length += 1
                if first_img:
                    img_p = os.path.join(self.path, img_f)
                    img = Image.open(img_p).convert('RGB')
                    self.width, self.height = img.size
                    first_img = False
        self.num_frames = length

    def __iter__(self):
        self.current_idx = 0  # maintain the index of current frame instead of getting property in CV2 to avoid bugs
        return self

    def __next__(self):
        # If all frames have been read at the beginning, raise StopIteration
        if self.current_idx == len(self):
            raise StopIteration

        self.current_idxs = []
        batch, times, indices = [], [], []

        num1 = 0
        for num in range(self.batch_size * self.step_size):
            if (num1 % self.step_size) == 0:
                self.current_idxs.append(num + 1 + self.current_idx)
            num1 += 1

        for current_id in self.current_idxs:
            if current_id >= len(self):
                break
            file_name = "74D401_%04d.jpg" % current_id  # 此处需要修改
            file_path = os.path.join(self.path, file_name)
            img = Image.open(file_path).convert('RGB')
            indices.append(self.current_idx)
            self.current_idx += self.step_size
            # timestamps_ms = (self.current_idx - 1) / self.fps * 1000  # FPS没有导入，暂时无法计算
            # times.append(timestamps_ms)
            if self.transform is not None:
                batch.append(self.transform(img))
            else:
                batch.append(img)
        if len(batch) == 0:
            raise StopIteration

        return batch, None, indices

    def __len__(self):
        return self.num_frames
